# -*- coding: utf-8 -*-
import base64
import logging
from os.path import join
import zlib

import msgpack
import requests

from messagebroker.json import json_dumps


logger = logging.getLogger(__name__)


class MessageBrokerClientV3(object):
    """
    V3版消息中间件协议

    文档参考：
    https://git.corp.qianka.com/chenlei/doc/wikis/message-broker-protocol#v3

    均使用中间件V3的JSON发送API
    """

    MARSHAL_JSON = 'json'
    MARSHAL_MSGPACK = 'msgpack'

    COMPRESS_ZLIB = 'zlib'

    def __init__(self):
        self.base_url = 'http://127.0.0.1:3000/message'
        self.timeout = 5

        self.session = requests.Session()
        # 可以控制连接池行为
        # adapter = requests.adapters.HTTPAdapter(pool_connections=10,
        #         pool_maxsize=10, max_retries=0, pool_block=False)
        # self.session.mount('http://', adapter)

    def configure(self, base_url, timeout):
        self.base_url = base_url
        self.timeout = timeout

    def send(
            self, connection, exchange, routing_key, o,
            marshal=None, compress=None,
            properties=None, headers=None):
        """
        发送Python对象，一般为dict类型
        会根据参数进行序列化、压缩等操作

        Args:
            :connection: str 发送AMQP集群目标
            :exchange: str AMQP消息交换站
            :routing_key: str AMQP消息路由键
            :o: object **可序列化的** 消息对象
            :marshal: MARSHAL_* 默认JSON编码
            :compress: COMPRSSS_* 默认不压缩
            :properties: AMQP标准消息属性，参考 0-9-1 协议文档
                         格式为map(string, string)
            :headers: 自定义头，格式为map(string, string)

        Returns:
            bool 是否发送到中间件
        """
        if marshal is None:
            marshal = self.MARSHAL_JSON

        if not self._check_marshal(marshal):
            raise RuntimeError('不支持的序列化模块: {}'.format(marshal))
        if compress and (not self._check_compress(compress)):
            raise RuntimeError('不支持的压缩模块: {}'.format(compress))

        if isinstance(o, bytes):
            # 发送内容为bytes时，转给sendb方法处理
            rv = self.sendb(
                connection, exchange, routing_key, o, properties, headers)
            return rv

        content_type = ''
        content_encoding = ''
        body = None

        if isinstance(o, str):
            # 发送内容为字符串时，不需要序列化
            content_type = 'text/plain'
            content_encoding = ''
            body = o
        else:
            # 其他情况按照参数进行序列化、压缩等处理
            # 序列化
            b, content_type = self._serialize(o, marshal)
            if compress:
                # 压缩
                b, content_encoding = self._compress(b, compress)

            # base64编码
            payload = base64.b64encode(b)
            body = payload.decode('ascii')

        rv = self.send_raw(
            connection, exchange, routing_key,
            content_type, content_encoding, body, properties, headers)
        return rv

    def sendb(
            self, connection, exchange, routing_key, b,
            properties=None, headers=None):
        """
        发送二进制内容，此类消息由发送者和接收者 **自行协商** 内容格式

        Args:
            :connection: str 发送AMQP集群目标
            :exchange: str AMQP消息交换站
            :routing_key: str AMQP消息路由键
            :b: bytes 二进制内容
            :properties: AMQP标准消息属性，参考 0-9-1 协议文档
                         格式为map(string, string)
            :headers: 自定义头，格式为map(string, string)

        Returns:
            bool 是否发送到中间件
        """

        if not isinstance(b, bytes):
            raise RuntimeError('参数b应该是bytes类型')

        content_type = ''
        content_encoding = 'binary'
        payload = base64.b64encode(b)
        body = payload.decode('ascii')

        return self.send_raw(
            connection, exchange, routing_key,
            content_type, content_encoding, body, properties, headers)

    def send_raw(
            self, connection, exchange, routing_key,
            content_type, content_encoding, body,
            properties=None, headers=None):
        """
        指定消息内容格式、编码格式，发送二进制内容的原始参数

        Args:
            :connection: str 发送AMQP集群目标
            :exchange: str AMQP消息交换站
            :routing_key: str AMQP消息路由键
            :content_type: str 消息内容格式
            :content_encoding: str 消息编码格式
            :body: str 消息内容
                   **注意**：如果是二进制内容需要提前base64编码
            :properties: AMQP标准消息属性，参考 0-9-1 协议文档
                         格式为map(string, string)
            :headers: 自定义头，格式为map(string, string)

        Returns:
            bool 是否发送到中间件
        """
        params = {
            'connection': connection,
            'exchange': exchange,
            'routing_key': routing_key,
            'content_type': content_type,
            'content_encoding': content_encoding,
            'body': body,
        }

        # AMQP消息属性与自定义指定的时候才添加参数
        if isinstance(properties, dict):
            params['properties'] = properties
        if isinstance(headers, dict):
            params['headers'] = headers

        url = join(self.base_url, 'v3/send')
        http_headers = {'Content-Type': 'application/json'}
        body = json_dumps(params)
        timeout = self.timeout

        rv = self._http_post(url, body, http_headers, timeout)
        return rv

    def _http_post(self, url, body, headers, timeout, retries=3):

        for i in range(retries):
            logger.debug('_http_post trial {}'.format(i))
            try:
                resp = self.session.post(
                        url, body, headers=headers, timeout=timeout)
                if resp.status_code == 200:
                    return True
            except requests.exceptions.RequestException as e:
                logger.exception('request error: {}'.format(e))
            except Exception as e:
                logger.exception('unknown error: {}'.format(e))

        logger.error('!!error after {} retries'.format(retries))

        return False

    def _check_marshal(self, marshal):
        """
        检查序列化模块是否合理

        Args:
            :marshal: MARSHAL_*

        Returns:
            bool 正常True 异常False
        """
        marshal_modeuls = (self.MARSHAL_JSON, self.MARSHAL_MSGPACK)
        return marshal in marshal_modeuls

    def _check_compress(self, compress):
        """
        检查压缩模块是否合理

        Args:
            :compress: COMPRESS_*

        Returns:
            bool 正常True 异常False
        """
        compress_modules = (self.COMPRESS_ZLIB,)
        return compress in compress_modules

    def _serialize(self, o, marshal):
        """
        序列化

        Args:
            :o: object
            :marshal: MARSHAL_*

        Returns:
            bytes, str 序列化后的二进制内容，以及 content_type
        """
        if marshal == self.MARSHAL_JSON:
            rv = json_dumps(o), 'application/json'
        elif marshal == self.MARSHAL_MSGPACK:
            rv = msgpack.packb(o), 'application/x-msgpack'
        else:
            raise RuntimeError('不支持的序列化模块：{}'.format(marshal))
        return rv

    def _compress(self, o, compress):
        """
        压缩

        Args:
            :o: object
            :compress: COMPRESS_*

        Returns:
            bytes, str 压缩吼的二进制内容，以及 content_encoding
        """
        if compress == self.COMPRESS_ZLIB:
            rv = zlib.compress(o, 1), 'zlib'
        else:
            raise RuntimeError('不支持的压缩模块：{}'.format(compress))
        return rv
